import json
from pathlib import Path
from docx import Document
from docx.shared import Pt
from bs4 import BeautifulSoup

TEXT_BLOCKS = {"text", "paragraph_title", "header", "footer"}  # 所有文字块类型


def html_table_to_docx_table(doc, html_content):
    soup = BeautifulSoup(html_content, "html.parser")
    table_tag = soup.find("table")
    if not table_tag:
        return

    rows = table_tag.find_all("tr")
    if not rows:
        return

    num_rows = len(rows)
    num_cols = len(rows[0].find_all(["td", "th"]))
    table = doc.add_table(rows=num_rows, cols=num_cols)
    table.style = 'Table Grid'

    for i, row in enumerate(rows):
        cells = row.find_all(["td", "th"])
        for j, cell in enumerate(cells):
            table.cell(i, j).text = cell.get_text(strip=True)
    doc.add_paragraph()  # 表格后空行


def export_to_word(json_paths, output_path):
    doc = Document()

    for idx, json_file in enumerate(json_paths):
        with open(json_file, 'r', encoding='utf-8') as f:
            data = json.load(f)

        doc.add_paragraph(f"--- Page {idx + 1} ---")

        for item in data.get('parsing_res_list', []):
            block_type = item.get('block_label')
            content = item.get('block_content')

            if block_type in TEXT_BLOCKS and content:
                p = doc.add_paragraph(content)
                p.style.font.size = Pt(11)
            elif block_type == 'table' and content:
                html_table_to_docx_table(doc, content)
            # seal 或其他类型可以根据需要单独处理

        if idx != len(json_paths) - 1:
            doc.add_page_break()

    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    doc.save(output_path)
    # print(f"Word 导出成功：{output_path}")
    return str(output_path)


# ------------------ 测试示例 ------------------
if __name__ == "__main__":
    # 假设我们有两个页面的结构化结果 JSON 文件
    json_files = [
        "output_structure/2f9f3436-91cd-466c-859a-40e08b7c77d9/page1.json",
        "output_structure/2f9f3436-91cd-466c-859a-40e08b7c77d9/page2.json",
        "output_structure/2f9f3436-91cd-466c-859a-40e08b7c77d9/page3.json",
        "output_structure/2f9f3436-91cd-466c-859a-40e08b7c77d9/page4.json"
    ]

    # 导出合并的 Word 文件
    export_to_word(json_files, "output_structure/result.docx")
